# encoding: utf-8

import cv2
import base64
import os
from glob import glob
import time
import aiohttp
import requests
import asyncio
import random
import numpy as np
import traceback
from typing import List, Tuple
from aiohttp import ClientConnectorError, ServerDisconnectedError

semaphore_limit = 2
sample_count = 5

# 去水印
# api = "http://182.92.239.24:6006/predict/rv_watermark"
# api = "https://u12453-8224-18ad9912.westb.seetacloud.com:8443/predict/rv_watermark"
api = "http://123.56.104.131:3389/predict/rv_watermark"

headers = {"Content-Type": "application/json"}
req_headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
}


async def fetch(session, img_url, name: str):
    start = time.time()
    response = requests.get(img_url, headers=req_headers)
    content = response.content
    b64_content = base64.b64encode(content).decode("utf-8")
    try:
        data = {"img": b64_content}
        async with session.post(
            url=api,
            headers=headers,
            json=data
        ) as response:
            end = time.time()
            # text = await response.text()
            # print(text)
            rs_data = await response.json()
            result_str = rs_data.get("results")
            content2 = base64.b64decode(result_str)
            # new_path = f"./images/predict/{name}"
            # with open(new_path, "wb") as f:
            #     f.write(content2)
            print(f"图片:{name}, 耗时：{end - start}")
        image1 = cv2.imdecode(np.frombuffer(content, np.uint8), cv2.IMREAD_COLOR)
        image2 = cv2.imdecode(np.frombuffer(content2, np.uint8), cv2.IMREAD_COLOR)
        combined_image = cv2.hconcat([image1, image2])
        cv2.imshow(name, combined_image)
        cv2.imwrite(name, combined_image)

    except Exception as e:
        print(traceback.format_exc())
        print(f"图片:{name}, 请求失败")


def get_random_img_list() -> List[Tuple[str, str]]:
    """

    :return: [file_abs_path, file_name]
    """
    temp = []
    path = r"E:/PythonCodes/fast_projects/fast-predict/images/test/images/origin/*.jpg"
    img_list = glob(path)
    files = random.sample(img_list, sample_count)
    # print(files)
    for file in files:
        base_name = os.path.basename(file)
        temp.append((file, base_name))

    return temp


def get_all_img_list() -> List[Tuple[str, str]]:
    temp = []
    path = r"E:/PythonCodes/fast_projects/fast-predict/images/test/images/origin/*.jpg"
    img_list = glob(path)
    for file in img_list:
        base_name = os.path.basename(file)
        temp.append((file, base_name))

    return temp


async def main(img_list: List[Tuple[str, str]]):
    async with aiohttp.ClientSession() as session:
        tasks = []
        semaphore = asyncio.Semaphore(semaphore_limit)
        for idx, (img_path, img_name) in enumerate(img_list):
            async def request(img_path, name):
                async with semaphore:
                    await fetch(session, img_path, name)

            tasks.append(request(img_path, img_name))

        await asyncio.gather(*tasks)
    cv2.waitKey(0)  # 等待任意按键
    cv2.destroyAllWindows()  # 销毁所有OpenCV窗口


if __name__ == '__main__':
    img_url = "https://private.zkh.com/PRODUCT/APP/PRO_AA2194018_03.jpg?x-oss-process=style/common_style_origin&timestamp=1695106223000"
    name = "demo.jpg"
    imgs = [(img_url, name)]
    asyncio.get_event_loop().run_until_complete(main(imgs))
